# -*- coding: utf-8 -*-

import pickle
from arch.package import NanoOpTypes
from arch.cxx_dumps import *

from ip.opdef import AllocOp, ExecOp, TranOp


def dump_pickle(op_struct: dict, filepath: str):
    # Use pickle directly dump the file.
    with open(filepath, 'wb') as pickle_file:
        pickle.dump(op_struct, pickle_file)


def dump_aop(aop: AllocOp):
    return cxx_dump_complex(aop.data) + \
           cxx_dump_int32(aop.addr) + \
           cxx_dump_int32(aop.op_code)


def dump_top(top: TranOp):
    return cxx_dump_int32(top.len_limit) + \
        cxx_dump_int32(top.op_code) + \
        cxx_dump_int32(top.dst_peer_id) + \
        cxx_dump_int32(top.dst_addr) + \
        cxx_dump_int32(top.dst_stride) + \
        cxx_dump_bool(top.dst_clear) + \
        cxx_dump_int32(top.dst_sb_addr) + \
        cxx_dump_int32(top.dst_ain_addr) + \
        cxx_dump_int32(top.src_addr) + \
        cxx_dump_int32(top.src_stride)


def dump_eop(eop: ExecOp):
    return cxx_dump_int32(eop.alu) + \
        cxx_dump_int32(eop.peer_id) + \
        cxx_dump_int32(eop.len_limit) + \
        cxx_dump_int32(eop.op_code) + \
        cxx_dump_bool(eop.dst_addr) + \
        cxx_dump_int32(eop.dst_increase) + \
        cxx_dump_bool(eop.dst_enable_offset) + \
        cxx_dump_int32(eop.dst_peer_id) + \
        cxx_dump_int32(eop.dst_stride) + \
        cxx_dump_bool(eop.dst_sb) + \
        cxx_dump_int32(eop.src_addr_0) + \
        cxx_dump_bool(eop.src_wait_passin_0) + \
        cxx_dump_bool(eop.src_wait_aluin_0) + \
        cxx_dump_bool(eop.src_enable_offset_0) + \
        cxx_dump_int32(eop.src_stride_0) + \
        cxx_dump_int32(eop.src_limit_0) + \
        cxx_dump_int32(eop.src_addr_1) + \
        cxx_dump_bool(eop.src_wait_passin_1) + \
        cxx_dump_bool(eop.src_wait_aluin_1) + \
        cxx_dump_bool(eop.src_enable_offset_1) + \
        cxx_dump_int32(eop.src_stride_1) + \
        cxx_dump_int32(eop.src_limit_1)


def dump_juna(op_struct: dict, filepath: str):
    # Format:
    # 4 bytes LE [Aop size] [Aops]
    # 4 bytes LE [Top size] [Tops]
    # 4 bytes LE [Eop size] [Eops]
    aops = op_struct[NanoOpTypes.AopsName]
    aops_data = cxx_dump_int32(len(aops))
    for op in aops:
        aops_data += dump_aop(op)
    tops = op_struct[NanoOpTypes.TopsName]
    tops_data = cxx_dump_int32(len(tops))
    for op in tops:
        tops_data += dump_top(op)
    eops = op_struct[NanoOpTypes.EopsName]
    eops_data = cxx_dump_int32(len(eops))
    for op in eops:
        eops_data += dump_eop(op)
    print("AllocOps: {}\nTranOps: {}\nExecOps: {}".format(
        len(aops), len(tops), len(eops)))
    # Output those data to file.
    with open(filepath, 'wb') as juna_file:
        juna_file.write(aops_data)
        juna_file.write(tops_data)
        juna_file.write(eops_data)


MARU_FORMATS = {
    'pickle': [dump_pickle, 'p'],
    'juna': [dump_juna, 'juna']
}